#include <iostream>
#include <iomanip>
#include <vector>
#include <array>
#include <string>
#include <cstring>
#include "Core/TensorExpr.h"
#include "FiniteDiff/LevelOp.h"
#include "FiniteDiff/FuncFiller.h"
#include "FiniteDiff/GhostFiller.h"
#include "FiniteDiff/MixedDer.h"
#include "FiniteDiff/Multigrid.h"
#include "../example/common/TestUtility.h"

#if DIM == 2

template class Tensor<Real,SpaceDim>;

const Vec<Real, 3> A { 1.0, 0.0, 2.0 };

const char *bcTypes = "DDDD";

class MixedDerLevelOp : public MixedDer, public MGLevelOp<2>
{
public:
  MixedDerLevelOp(const RectDomain<2> &aDomain, const Vec<Real, 3> &aA, Real w4j, const char *aBcTypes)
      : MixedDer(aDomain, aA), MGLevelOp(aDomain, w4j, aBcTypes)
  {
  }

  ~MixedDerLevelOp() = default;

  void applySmoother(const Tensor<Real, Dim> &phi, Tensor<Real, Dim> &smoothed, const Tensor<Real, Dim> &rhs) const
  {
    MixedDer::relax(phi, rhs, smoothed, weightForJacobi);
  }

  void computeResidual(const Tensor<Real, Dim> &phi, Tensor<Real, Dim> &rsd, const Tensor<Real, Dim> &rhs) const
  {
    MixedDer::apply(phi, rhs, rsd);
  }
};

template <int Dim, class T1, class T2>
std::array<Real,3> testMixedDer(const std::vector<const MGLevelOp<Dim> *> &opHier,
                                const std::vector<const Intergrid<Dim> *> &trHier,
                                const T1 &PHI, const T2 &D2PHI)
{
  const RectDomain<Dim> &rd = opHier[0]->getDomain();
  LevelOp<Dim> lvop(rd);
  FuncFiller<Dim> ffill(rd);
  Box<Dim> gbx = rd.getGhostedBox();
  Tensor<Real,Dim> phi(gbx);
  Tensor<Real,Dim> err(gbx);
  Tensor<Real,Dim> rhs(gbx);
  Tensor<Real,Dim-1> cData[2*Dim];

  for(int d = 0; d < Dim; ++d)
    for(int k : {0, 1})
      cData[2*d+k].resize(reduce(gbx, d));
  ffill.fillBdryValues(cData, CellCentered, PHI);
  ffill.fillPointValues(err, CellCentered, PHI);
  ffill.fillPointValues(rhs, CellCentered, D2PHI);
  phi = 0.0;

  Multigrid<Dim> solver(opHier, trHier);
  solver.setParam({2, 2, 20}, 30, 1.001, 1e-13);
  solver.solve(phi, rhs, cData);

  err = err - phi;
  std::array<Real, 3> errnorm;
  for(int p=0; p<=2; ++p)
    errnorm[p] = lvop.NRM(err, p);
  return errnorm;
}

template <int Dim, class T1, class T2>
void doTest(const std::vector<int> &gridSize, int numCompHier, const T1 &PHI, const T2 &D2PHI)
{
  const int numGrid = gridSize.size();
  std::vector<const RectDomain<Dim> *> gridHier;
  std::vector<const MGLevelOp<Dim> *> opHier;
  std::vector<const Intergrid<Dim> *> trHier;
  for(int n = 0; n < numGrid; ++n) {
    const int N = gridSize[n];
    const RectDomain<Dim> *pRd = new RectDomain<Dim>(Box<Dim>(0, N-1), 1.0/N, CellCentered, 2);
    gridHier.push_back(pRd);
    const MGLevelOp<Dim> *plv = new MixedDerLevelOp(*pRd, A, 0.6, bcTypes);
    opHier.push_back(plv);
    if(n != 0) {
      const Intergrid<Dim> *pt = new Intergrid<Dim>(*pRd, *(gridHier[n-1]));
      trHier.push_back(pt);
    }
  }

  std::vector<std::array<Real, 3>> errnorm(numCompHier);
  for(int n = numCompHier-1; n >= 0; --n) {
//    std::vector<const RectDomain<Dim> *> _gridHier(gridHier.cbegin() + n, gridHier.cend());
    std::vector<const MGLevelOp<Dim> *> _opHier(opHier.cbegin() + n, opHier.cend());
    std::vector<const Intergrid<Dim> *> _trHier(trHier.cbegin() + n, trHier.cend());
    errnorm[n] = testMixedDer<Dim>(_opHier, _trHier, PHI, D2PHI);
  }

  for(int n = numGrid-1; n >= 0; --n) {
    if(n != 0)
      delete trHier[n-1];
    delete opHier[n];
    delete gridHier[n];
  }
  printConvergenceTable(&gridSize[0], errnorm);
}

int main()
{
  reset_dbglevel(1);
  std::cout << "==============================================================================" <<  std::endl;
  doTest<2>({512, 256, 128, 64, 32, 16, 8},
            3,
            [](Vec<Real, 2> x) {
              return sin(2 * M_PI * x[0]) * cos(2 * M_PI * x[1]);
            },
            [&](Vec<Real, 2> x) {
              Real xx = (-4*M_PI*M_PI) * sin(2*M_PI*x[0]) * cos(2*M_PI*x[1]);
              Real xy = (-4*M_PI*M_PI) * cos(2*M_PI*x[0]) * sin(2*M_PI*x[1]);
              Real yy = (-4*M_PI*M_PI) * sin(2*M_PI*x[0]) * cos(2*M_PI*x[1]);
              return A[0]*xx + A[1]*xy + A[2]*yy;
            });
  std::cout << std::endl;
  return 0;
}

#else

#warning "testMG4MixedDer is omitted for 3D. "
int main()
{
  return 0;
}

#endif // DIM == 2
